# 要求：tensorflow >= 1.8
# toco --output_file=/hadoop/station_photo_determine/model/newmodel/tflite_model/ws3_vgg_model_1106.tflite --output_format=TFLITE --keras_model_file=/hadoop/station_photo_determine/model/newmodel/ws3_vgg_model_1106.h5
# 或者
# tflite_convert --keras_model_file=/hadoop/station_photo_determine/model/newmodel/ws3_vgg_model_1106.h5 --output_file=/station_photo_determine/model/newmodel/tflite_model/ws3_vgg_model_1106.tflite --output_format=tflite

# 以下或者更容易定制化
from keras.backend import clear_session
import numpy as np
import tensorflow as tf
from tensorflow.contrib import lite
clear_session()
np.set_printoptions(suppress=True)
input_graph_name = "/aidfs/003/model_bak/LabelCheck_20190905_1906.h5"
output_graph_name = input_graph_name[:-3] + '.tflite'
converter = lite.TFLiteConverter.from_keras_model_file(model_file=input_graph_name)
converter.post_training_quantize = True # 量化模型
#在windows平台这个函数有问题，无法正常使用
tflite_model = converter.convert()
open(output_graph_name, "wb").write(tflite_model)
print ("generate:",output_graph_name)
# 如果有自定义对象:
"""
from tensorflow.python.keras.utils import CustomObjectScope
def relu6(x):
    return K.relu(x, max_value=6)
with CustomObjectScope({'relu6': relu6}):
    tflite_model = tf.contrib.lite.TFLiteConverter.from_keras_model_file('tflite-models/keras_model_converted.h5').convert()
    with open('tflite-models/model.tflite', 'wb') as f:
        f.write(tflite_model)
"""